package median

import (
	"sort"
	"gitee.com/tfcolin/dsg"
)

type MedianData struct {
	n int
	z []float64
	omega []float64
	w float64
}

func linear_median (v []float64) float64 {
	v1 := make ([]float64, len(v))
	copy (v1, v)
	sort.Float64s (v1)
	return v1[(len(v1) - 1) / 2]
}

func (md MedianData) Len () int {
	return md.n
}

func (md * MedianData) Swap (i, j int) {
	md.z[i], md.z[j] = md.z[j], md.z[i]
	md.omega[i], md.omega[j] = md.omega[j], md.omega[i]
}

func (md MedianData) Less (i, j int) bool {
	if md.z[i] < md.z[j] {
		return true
	} else {
		return false
	}
}

func MedianSlow (md * MedianData) float64 {
	sort.Sort (md)
	sum := 0.0
	for i := 0; i < md.n; i ++ {
		sum += md.omega[i]
		if sum >= md.w {
			return md.z[i]
		}
	}
	return 0
}

func Median (md MedianData) float64 {
	n := md.n
	n5 := (n - 1) / 5 + 1
	zm := make ([]float64, n5)

	for i := 0; i < n5; i ++ {
		zin := md.z[i * 5 : dsg.IMin ((i + 1) * 5, n)]
		zm[i]= linear_median (zin)
	}

	var zmm float64
	if n5 <= 5 {
		zmm = linear_median (zm)
	} else {
		md1 := MedianData {
			n: n5,
			z: zm,
			omega : make ([]float64, n5),
			w : float64(n5) * .5,
		}
		for i := 0; i < n5; i ++ {
			md1.omega[i] = 1
		}
		zmm = Median (md1)
	}

	var sum_small, sum_equal float64
	var n_small, n_equal, n_big int
	var z_small, z_big, omega_small, omega_big []float64

	for i := 0; i < n; i ++ {
		switch {
		case md.z[i] < zmm: 
			sum_small += md.omega[i]
			n_small ++
		case md.z[i] == zmm:
			sum_equal +=  md.omega[i]
			n_equal ++
		}
	}
	n_big = n - n_small - n_equal

	z_small = make ([]float64, n_small)
	omega_small = make ([]float64, n_small)
	z_big = make ([]float64, n_big)
	omega_big = make ([]float64, n_big)

	is, ib := 0, 0
	for i := 0; i < n; i ++ {
		switch {
		case md.z[i] < zmm: 
			z_small[is] = md.z[i]
			omega_small[is] = md.omega[i]
			is ++
		case md.z[i] > zmm:
			z_big[ib] = md.z[i]
			omega_big[ib] = md.omega[i]
			ib ++
		}
	}

	switch {
	case sum_small + sum_equal < md.w :
		md2 := MedianData {
			n_big,
			z_big,
			omega_big,
			md.w - (sum_small + sum_equal),
		}
		return Median (md2)
	case md.w <= sum_small :
		md2 := MedianData {
			n_small,
			z_small,
			omega_small,
			md.w,
		}
		return Median (md2)
	default: 
		return zmm
	}
}
